#!/usr/bin/env python
"""
Command line utility to calculate a position specific scoring matrix
for any set of sequences.

Writes tab-delimited table with residue letters in the header line
and in each column the probability of observing each of those residues
calculated from the sequence file.

Example:-
 $ python bin/ACGTCounts.py [insequencefile].fas > PSM.txt
 or
 $ python ACGTCounts.py [insequencefile].fas -out PSM.txt


"""
# Copyright (c) Alex Leach. 2011
# Contact: Alex Leach (albl500@york.ac.uk)
# University of York, Department of Biology,
# Wentworth Way, York YO10 4DU, United Kingdom
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have receive a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>
#
# Description:
# Calculates position specific scoring matrix from as sequence file.


import sys
from Bio import SeqIO
from ssummo.ssummolib import Options

def posCounts( inFile, fileFormat='fasta' ):
        countDict = {}
        iterator = SeqIO.parse( inFile, format=fileFormat  )
        resSeen = set()
        nSeqs = 0
        for seqRec in iterator:
                pos = 0
                nSeqs += 1
                for res in seqRec.seq.tostring():
                        res = res.upper()
                        if pos not in countDict.keys():
                                countDict.update( { pos : { res : 1 } } )
                        elif res in countDict[pos].keys():
                                countDict[pos][res] += 1
                        else:
                                countDict[pos].update( { res : 1 } )
                        resSeen.add(res)
                        pos += 1
        resSeen = sorted( list( resSeen ) )
        return countDict,resSeen,nSeqs

def tabularise( PSMDict ,nseqs, out_file, residues_seen = []):
        if len(residues_seen) == 0:
                residues = sorted( list( set( PSMDict.values().keys() )) )   # residues all unique, and in order
        else:
                residues = residues_seen
        if '-' in residues:
                del( residues[ residues.index( '-' ) ] )
        sys.stdout.write( '\t'.join( residues ) + '\n' )
        for i in xrange( len( PSMDict.keys() ) ):
                outstr = ''
                if '-' in PSMDict[i].keys():
                        blankCount = PSMDict[i]['-']
                else:
                        blankCount = 0
                for residue in residues:
                        if residue in PSMDict[i].keys():
                                outstr += str( float(PSMDict[i][residue]) / float(nseqs - blankCount) ) + '\t'
                        else:
                                outstr += '0\t'
                out_file.write( outstr.rstrip() + '\n' )

def parseArgs( args ):
    options = Options()
    options.options = { '-format' : 'fasta' , '-in' : sys.stdin ,'-out':sys.stdout}
    options.helpTxt = { '-format' : 'Input sequence format. [fasta]',
                        '-in'     : 'Input sequence file. [/dev/stdin]',
                        '-out'    : 'Output file for PSM. [/dev/stdout]'
                        }
    options.multiargs = []
    options.singleargs = ['-in','-format','-out']
    options.parseArgs( args )
    return options

if __name__ == '__main__':
        args = sys.argv[1:]
        options = parseArgs( args )
        if options['-in'] != sys.stdin:
                options[ '-in' ] = file( options['-in'],'r')
        if options['-out'] != sys.stdout:
                options[ '-out' ] = file( options['-out'] ,'w' ) 
        countDict, seen,nSeqs = posCounts( options['-in'],fileFormat = options['-format'])
        tabularise( countDict , nSeqs,options['-out'], residues_seen = seen )
        options['-in'].close()
        options['-out'].close()
